import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

def distance(a, b):
    return np.sqrt(np.sum((a - b) ** 2))

class KNN:
    def __init__(self, k, label_num):
        self.k = k
        self.label_num = label_num  

    def fit(self, x_train, y_train):
        self.x_train = x_train
        self.y_train = y_train

    def get_knn_indices(self, x):
        dis = list(map(lambda a: distance(a, x), self.x_train))
        knn_indices = np.argsort(dis)[:self.k]
        return knn_indices

    def get_label(self, x):
        knn_indices = self.get_knn_indices(x)
        label_statistic = np.zeros(self.label_num)
        for idx in knn_indices:
            label_statistic[int(self.y_train[idx])] += 1
        return np.argmax(label_statistic)

    def predict(self, x_test):
        preds = np.zeros(len(x_test), dtype=int)
        for i, x in enumerate(x_test):
            preds[i] = self.get_label(x)
        return preds

np.random.seed(0)
x0 = np.random.randn(50, 2) + np.array([2, 2])
y0 = np.zeros(50)
x1 = np.random.randn(50, 2) + np.array([7, 7])
y1 = np.ones(50)

x_train = np.vstack([x0, x1])
y_train = np.hstack([y0, y1])


plt.figure()
plt.scatter(x_train[y_train==0,0], x_train[y_train==0,1], c='blue', marker='o', label='Class 0')
plt.scatter(x_train[y_train==1,0], x_train[y_train==1,1], c='red', marker='x', label='Class 1')
plt.xlabel('X')
plt.ylabel('Y')
plt.title('Simple 2D Dataset')
plt.legend()
plt.show()

np.random.seed(1)  
k = np.random.choice(np.arange(1, 11))  
print(f"随机选择的 K 值为: {k}")

cmap_light = ListedColormap(['lightblue', 'lightcoral'])
x_min, x_max = x_train[:,0].min()-1, x_train[:,0].max()+1
y_min, y_max = x_train[:,1].min()-1, x_train[:,1].max()+1
xx, yy = np.meshgrid(np.arange(x_min, x_max, 0.1),
                     np.arange(y_min, y_max, 0.1))
grid_data = np.c_[xx.ravel(), yy.ravel()]

knn = KNN(k=k, label_num=2)
knn.fit(x_train, y_train)
z = knn.predict(grid_data)

plt.figure(figsize=(7,6))
plt.pcolormesh(xx, yy, z.reshape(xx.shape), cmap=cmap_light, alpha=0.5)
plt.scatter(x_train[y_train==0,0], x_train[y_train==0,1], c='blue', marker='o', label='Class 0')
plt.scatter(x_train[y_train==1,0], x_train[y_train==1,1], c='red', marker='x', label='Class 1')
plt.xlabel('X')
plt.ylabel('Y')
plt.title(f'KNN Classification (K={k})')
plt.legend()
plt.show()

pred_train = knn.predict(x_train)
accuracy = np.mean(pred_train == y_train)
print(f"K = {k} 时，训练集准确率 = {accuracy*100:.2f}%")